import os
import re
import shutil
import random
from collections import defaultdict

import numpy as np


# def split_dataset(source_dir, target_root, ratio=(0.8, 0.2), seed=42):
#     # 检查源文件夹是否存在
#     if not os.path.exists(source_dir):
#         print(f"源文件夹 {source_dir} 不存在。")
#         return
#     # 按前缀分组图片
#     prefix_groups = defaultdict(list)
#     for filename in os.listdir(source_dir):
#         prefix = filename.split('_')[0]
#         file_path = os.path.join(source_dir, filename)
#         prefix_groups[prefix].append((filename, file_path))
#
#     # 创建目标文件夹
#     train_dir = os.path.join(target_root, 'train')
#     val_dir = os.path.join(target_root, 'val')
#     os.makedirs(train_dir, exist_ok=True)
#     os.makedirs(val_dir, exist_ok=True)
#
#     # 对每个分组按比例划分并复制文件
#     random.seed(seed)
#     for prefix, files in prefix_groups.items():
#         total = len(files)
#         if total == 0:
#             continue
#         train_num = int(total * ratio[0])
#         val_num = total - train_num
#
#         random.shuffle(files)
#         train_files = files[:train_num]
#         val_files = files[train_num:]
#
#         def copy_files(file_list, target_dir):
#             for filename, src_path in file_list:
#                 target_file_path = os.path.join(target_dir, filename)
#                 try:
#                     shutil.move(src_path, target_file_path)
#                 except Exception as e:
#                     print(f"复制 {filename} 时出错: {e}")
#
#         copy_files(train_files, train_dir)
#         copy_files(val_files, val_dir)
#         print(f"前缀 {prefix}:共 {total} 张，划分 train:{len(train_files)} | val:{len(val_files)}")

def split_dataset(source_dir, target_root, ratio=(0.8, 0.2), seed=42):
    # 检查源文件夹是否存在
    if not os.path.exists(source_dir):
        print(f"源文件夹 {source_dir} 不存在。")
        return

    # 按前缀分组图片
    prefix_groups = defaultdict(list)
    all_classes = set()  # 用于收集所有类别

    for filename in os.listdir(source_dir):
        # 提取类别（假设文件名格式为'class_数字_...'）
        class_name = filename.split('_')[0]
        all_classes.add(class_name)

        file_path = os.path.join(source_dir, filename)
        prefix_groups[class_name].append((filename, file_path))

    # 创建目标文件夹
    train_dir = os.path.join(target_root, 'train')
    val_dir = os.path.join(target_root, 'val')
    os.makedirs(train_dir, exist_ok=True)
    os.makedirs(val_dir, exist_ok=True)

    # 生成类别映射
    classes = sorted(all_classes)
    classes_to_idx = {cls: i for i, cls in enumerate(classes)}
    idx_to_labels = {i: cls for i, cls in enumerate(classes)}

    # 保存映射为txt文件
    mapping_dir = target_root  # 映射文件保存在目标根目录

    # 保存classes_to_idx
    with open(os.path.join(mapping_dir, 'classes_to_label.txt'), 'w', encoding='utf-8') as f:
        for cls, idx in classes_to_idx.items():
            f.write(f"{cls},{idx}\n")

    # 保存idx_to_labels
    with open(os.path.join(mapping_dir, 'label_to_classes.txt'), 'w', encoding='utf-8') as f:
        for idx, cls in idx_to_labels.items():
            f.write(f"{idx},{cls}\n")

    # 保存类别列表到classes.txt
    with open(os.path.join(mapping_dir, 'classes.txt'), 'w', encoding='utf-8') as f:
        for cls in classes:
            f.write(f"{cls}\n")

    print(f"类别映射已保存至 {mapping_dir}")
    print(f"类别列表已保存至 {os.path.join(mapping_dir, 'classes.txt')}")

    # 对每个分组按比例划分并复制文件
    random.seed(seed)
    for class_name, files in prefix_groups.items():
        total = len(files)
        if total == 0:
            continue

        train_num = int(total * ratio[0])
        val_num = total - train_num

        random.shuffle(files)
        train_files = files[:train_num]
        val_files = files[train_num:]

        def copy_files(file_list, target_dir):
            for filename, src_path in file_list:
                target_file_path = os.path.join(target_dir, filename)
                try:
                    shutil.move(src_path, target_file_path)
                except Exception as e:
                    print(f"移动 {filename} 到 {target_dir} 时出错: {e}")

        copy_files(train_files, train_dir)
        copy_files(val_files, val_dir)
        print(f"类别 {class_name}: 共 {total} 张，划分 train: {len(train_files)} | val: {len(val_files)}")


# def split_dataset_flat(source_dir, target_root, ratio=(0.8, 0.2), seed=42):
#     # 检查源文件夹是否存在
#     if not os.path.exists(source_dir):
#         print(f"源文件夹 {source_dir} 不存在。")
#         return
#     # 检查比例合法性
#     if round(sum(ratio), 2) != 1.0:
#         print("划分比例总和必须为1（如(0.8,0.2)），请重新设置ratio参数。")
#         return
#
#     # 1. 按类别分组图片（同时提取数字后缀）
#     prefix_groups = defaultdict(list)  # key: 类别名, value: [(原文件名, 数字后缀, 文件路径), ...]
#     all_classes = set()
#     img_suffixes = ('.jpg', '.jpeg', '.png', '.bmp', '.tiff')  # 支持的图片后缀
#     # 正则表达式:匹配文件名中的数字后缀（如"cat_001.jpg"中的"001"）
#     # 模式说明:\d+ 匹配1个以上数字，\.\w+$ 匹配扩展名（如.jpg）
#     number_pattern = re.compile(r'(\d+)\.\w+$')
#
#     for filename in os.listdir(source_dir):
#         # 过滤非图片文件
#         if not filename.lower().endswith(img_suffixes):
#             print(f"跳过非图片文件:{filename}")
#             continue
#
#         # 提取类别（假设格式为"类别_数字.后缀"，如"cat_001.jpg"）
#         class_name = filename.split('_')[0]  # 按'_'分割取前缀作为类别
#         all_classes.add(class_name)  # 【关键修复】将类别名添加到all_classes集合
#
#         # 提取数字后缀（核心:从原文件名中获取数字部分）
#         match = number_pattern.search(filename)
#         if not match:
#             print(f"警告:文件名 {filename} 中未找到数字后缀，将跳过该文件")
#             continue
#         number_suffix = match.group(1)  # 提取数字部分（如"001"）
#
#         file_path = os.path.join(source_dir, filename)
#         prefix_groups[class_name].append((filename, number_suffix, file_path))
#
#     if not all_classes:
#         print("源文件夹中未找到有效图片文件，无法划分数据集。")
#         return
#     print(f"共检测到 {len(all_classes)} 个类别，开始划分数据集...")
#
#     # 2. 创建目标文件夹结构（train/val → 类别子文件夹）
#     train_root = os.path.join(target_root, 'train')
#     val_root = os.path.join(target_root, 'val')
#     os.makedirs(train_root, exist_ok=True)
#     os.makedirs(val_root, exist_ok=True)
#
#     # 为每个类别创建train/val子文件夹
#     for class_name in all_classes:
#         os.makedirs(os.path.join(train_root, class_name), exist_ok=True)
#         os.makedirs(os.path.join(val_root, class_name), exist_ok=True)
#     print(f"目标文件夹结构已创建:\n- {train_root}\n- {val_root}")
#
#     # 3. 保存类别映射文件
#     classes = sorted(all_classes)
#     classes_to_idx = {cls: i for i, cls in enumerate(classes)}
#     idx_to_labels = {i: cls for i, cls in enumerate(classes)}
#
#     mapping_dir = target_root
#     with open(os.path.join(mapping_dir, 'classes_to_label.txt'), 'w', encoding='utf-8') as f:
#         for cls, idx in classes_to_idx.items():
#             f.write(f"{cls},{idx}\n")
#     with open(os.path.join(mapping_dir, 'label_to_classes.txt'), 'w', encoding='utf-8') as f:
#         for idx, cls in idx_to_labels.items():
#             f.write(f"{idx},{cls}\n")
#     with open(os.path.join(mapping_dir, 'classes.txt'), 'w', encoding='utf-8') as f:
#         for cls in classes:
#             f.write(f"{cls}\n")
#     print(f"类别映射文件已保存至 {mapping_dir}")
#
#     # 4. 按类别划分并移动图片（核心:使用原数字后缀重命名）
#     random.seed(seed)
#     for class_name, files in prefix_groups.items():
#         total_imgs = len(files)
#         if total_imgs == 0:
#             print(f"类别 {class_name} 无有效图片，跳过划分。")
#             continue
#
#         # 计算划分数量
#         train_num = int(total_imgs * ratio[0])
#         val_num = total_imgs - train_num
#
#         # 随机打乱并划分
#         random.shuffle(files)
#         train_files = files[:train_num]
#         val_files = files[train_num:]
#
#         # 定义带重命名的移动函数（使用原数字后缀）
#         def move_with_original_suffix(file_list, target_root_dir, cls_name):
#             """
#             将文件移动到目标文件夹，新文件名 = 原数字后缀 + 原扩展名
#             例如:原文件cat_001.jpg → 目标文件train/cat/001.jpg
#             """
#             target_class_dir = os.path.join(target_root_dir, cls_name)
#             for filename, number_suffix, src_path in file_list:
#                 # 提取原文件扩展名（如.jpg、.png）
#                 _, ext = os.path.splitext(filename)
#                 # 新文件名:数字后缀 + 扩展名（如001.jpg）
#                 new_filename = f"{number_suffix}{ext}"
#                 target_file_path = os.path.join(target_class_dir, new_filename)
#
#                 # 处理可能的文件名重复（若同一类别下有相同数字后缀的文件）
#                 if os.path.exists(target_file_path):
#                     print(f"警告:{new_filename} 已存在，将重命名为 {number_suffix}_dup{ext}")
#                     new_filename = f"{number_suffix}_dup{ext}"
#                     target_file_path = os.path.join(target_class_dir, new_filename)
#
#                 try:
#                     shutil.move(src_path, target_file_path)
#                 except Exception as e:
#                     print(f"处理 {filename} 时出错:{str(e)}")
#
#         # 移动并命名train和val图片
#         move_with_original_suffix(train_files, train_root, class_name)
#         move_with_original_suffix(val_files, val_root, class_name)
#
#         # 打印划分结果
#         print(f"类别 {class_name:10} | 总数: {total_imgs:4} | train: {len(train_files):4} | val: {len(val_files):4}")
#
#     print(f"\n数据集划分完成！最终结构:")
#     print(f"- 训练集:{train_root}/[类别名]/[原数字后缀].jpg（如001.jpg）")
#     print(f"- 验证集:{val_root}/[类别名]/[原数字后缀].jpg（如101.jpg）")
#     print(f"- 类别映射:{mapping_dir}/[classes_to_label.txt等]")

# def split_dataset_flat(source_dir, target_root, ratio=(0.7, 0.15, 0.15), seed=42):
#     # 检查源文件夹是否存在
#     if not os.path.exists(source_dir):
#         print(f"源文件夹 {source_dir} 不存在。")
#         return
#
#     # 检查比例合法性：必须是2或3个元素，且总和为1
#     if len(ratio) not in (2, 3):
#         print("ratio参数必须包含2个或3个元素（如(0.8,0.2)或(0.7,0.2,0.1)）")
#         return
#     if round(sum(ratio), 2) != 1.0:
#         print("划分比例总和必须为1（如(0.8,0.2)或(0.7,0.2,0.1)），请重新设置ratio参数。")
#         return
#
#     # 确定数据集类型（2个比例：train+val；3个比例：train+val+test）
#     dataset_types = ['train', 'val'] if len(ratio) == 2 else ['train', 'val', 'test']
#     print(f"将划分数据集为: {dataset_types}，比例: {ratio}")
#
#     # 1. 按类别分组图片（同时提取数字后缀）
#     prefix_groups = defaultdict(list)  # key: 类别名, value: [(原文件名, 数字后缀, 文件路径), ...]
#     all_classes = set()
#     img_suffixes = ('.jpg', '.jpeg', '.png', '.bmp', '.tiff')  # 支持的图片后缀
#     # 正则表达式:匹配文件名中的数字后缀（如"cat_001.jpg"中的"001"）
#     number_pattern = re.compile(r'(\d+)\.\w+$')
#
#     for filename in os.listdir(source_dir):
#         # 过滤非图片文件
#         if not filename.lower().endswith(img_suffixes):
#             print(f"跳过非图片文件:{filename}")
#             continue
#
#         # 提取类别（假设格式为"类别_数字.后缀"，如"cat_001.jpg"）
#         class_name = filename.split('_')[0]  # 按'_'分割取前缀作为类别
#         all_classes.add(class_name)
#
#         # 提取数字后缀
#         match = number_pattern.search(filename)
#         if not match:
#             print(f"警告:文件名 {filename} 中未找到数字后缀，将跳过该文件")
#             continue
#         number_suffix = match.group(1)  # 提取数字部分（如"001"）
#
#         file_path = os.path.join(source_dir, filename)
#         prefix_groups[class_name].append((filename, number_suffix, file_path))
#
#     if not all_classes:
#         print("源文件夹中未找到有效图片文件，无法划分数据集。")
#         return
#     print(f"共检测到 {len(all_classes)} 个类别，开始划分数据集...")
#
#     # 2. 创建目标文件夹结构（动态创建所有需要的数据集目录）
#     dataset_roots = {
#         dtype: os.path.join(target_root, dtype)
#         for dtype in dataset_types
#     }
#
#     # 创建根目录
#     for root in dataset_roots.values():
#         os.makedirs(root, exist_ok=True)
#
#     # 为每个类别创建子文件夹
#     for class_name in all_classes:
#         for root in dataset_roots.values():
#             os.makedirs(os.path.join(root, class_name), exist_ok=True)
#
#     print(f"目标文件夹结构已创建:")
#     for dtype, root in dataset_roots.items():
#         print(f"- {dtype}: {root}")
#
#     # 3. 保存类别映射文件
#     classes = sorted(all_classes)
#     classes_to_idx = {cls: i for i, cls in enumerate(classes)}
#     idx_to_labels = {i: cls for i, cls in enumerate(classes)}
#
#     mapping_dir = target_root
#     with open(os.path.join(mapping_dir, 'classes_to_label.txt'), 'w', encoding='utf-8') as f:
#         for cls, idx in classes_to_idx.items():
#             f.write(f"{cls},{idx}\n")
#     with open(os.path.join(mapping_dir, 'label_to_classes.txt'), 'w', encoding='utf-8') as f:
#         for idx, cls in idx_to_labels.items():
#             f.write(f"{idx},{cls}\n")
#     with open(os.path.join(mapping_dir, 'classes.txt'), 'w', encoding='utf-8') as f:
#         for cls in classes:
#             f.write(f"{cls}\n")
#     print(f"类别映射文件已保存至 {mapping_dir}")
#
#     # 4. 按类别划分并移动图片（核心:使用原数字后缀重命名，确保val和test为偶数）
#     random.seed(seed)
#     for class_name, files in prefix_groups.items():
#         total_imgs = len(files)
#         if total_imgs == 0:
#             print(f"类别 {class_name} 无有效图片，跳过划分。")
#             continue
#
#         # 计算各数据集的数量（确保val和test为偶数）
#         if len(ratio) == 2:
#             # 处理 train + val 模式（仅确保val为偶数）
#             train_num = int(total_imgs * ratio[0])
#             val_num = total_imgs - train_num
#
#             # 确保val_num为偶数
#             if val_num % 2 != 0:
#                 val_num += 1 if val_num > 0 else 1  # 避免0变成负数
#                 train_num = total_imgs - val_num  # 重新计算train_num
#                 if train_num < 0:  # 极端情况处理（总样本1时）
#                     train_num = 0
#                     val_num = total_imgs if total_imgs % 2 == 0 else 0
#
#             dataset_nums = [train_num, val_num]
#         else:
#             # 处理 train + val + test 模式（确保val和test为偶数）
#             train_num = int(total_imgs * ratio[0])
#             val_num = int(total_imgs * ratio[1])
#             test_num = total_imgs - train_num - val_num
#
#             # 先确保val_num为偶数
#             if val_num % 2 != 0:
#                 val_num += 1 if val_num > 0 else 1
#                 test_num = total_imgs - train_num - val_num  # 调整test_num
#
#             # 再确保test_num为偶数
#             if test_num % 2 != 0:
#                 test_num += 1 if test_num > 0 else 1
#                 train_num = total_imgs - val_num - test_num  # 调整train_num
#
#             # 最终安全检查（避免train_num为负）
#             if train_num < 0:
#                 train_num = 0
#                 # 重新分配val和test为偶数且总和为total_imgs
#                 val_num = total_imgs // 2 if total_imgs % 2 == 0 else (total_imgs // 2)
#                 val_num = val_num if val_num % 2 == 0 else val_num - 1
#                 test_num = total_imgs - val_num
#
#             dataset_nums = [train_num, val_num, test_num]
#
#         # 随机打乱并划分
#         random.shuffle(files)
#         dataset_files = []
#         start = 0
#         for num in dataset_nums:
#             end = start + num
#             dataset_files.append(files[start:end])
#             start = end
#
#         # 定义带重命名的移动函数（使用原数字后缀）
#         def move_with_original_suffix(file_list, target_root_dir, cls_name):
#             """将文件移动到目标文件夹，新文件名 = 原数字后缀 + 原扩展名"""
#             target_class_dir = os.path.join(target_root_dir, cls_name)
#             for filename, number_suffix, src_path in file_list:
#                 # 提取原文件扩展名
#                 _, ext = os.path.splitext(filename)
#                 # 新文件名:数字后缀 + 扩展名（如001.jpg）
#                 new_filename = f"{number_suffix}{ext}"
#                 target_file_path = os.path.join(target_class_dir, new_filename)
#
#                 # 处理可能的文件名重复
#                 if os.path.exists(target_file_path):
#                     print(f"警告:{new_filename} 已存在，将重命名为 {number_suffix}_dup{ext}")
#                     new_filename = f"{number_suffix}_dup{ext}"
#                     target_file_path = os.path.join(target_class_dir, new_filename)
#
#                 try:
#                     shutil.copy2(src_path, target_file_path)
#                 except Exception as e:
#                     print(f"处理 {filename} 时出错:{str(e)}")
#
#         # 移动并命名各数据集图片
#         for i, dtype in enumerate(dataset_types):
#             move_with_original_suffix(dataset_files[i], dataset_roots[dtype], class_name)
#
#         # 打印划分结果（标注偶数信息）
#         info_parts = []
#         for dtype, files in zip(dataset_types, dataset_files):
#             count = len(files)
#             even_note = "(偶数)" if (dtype in ['val', 'test'] and count % 2 == 0) else ""
#             info_parts.append(f"{dtype}: {count:4} {even_note}")
#
#         print(f"类别 {class_name:10} | 总数: {total_imgs:4} | " + " | ".join(info_parts))
#
#     print(f"\n数据集划分完成！最终结构:")
#     for dtype, root in dataset_roots.items():
#         print(f"- {dtype}集:{root}/[类别名]/[原数字后缀].jpg（如001.jpg）")
#     print(f"- 类别映射:{mapping_dir}/[classes_to_label.txt等]")

def split_dataset_flat(source_dir, target_root, ratio=(0.7, 0.15, 0.15), seed=42):
    """
    严格按照指定比例划分数据集，确保val和test为偶数且保持比例关系

    参数:
        source_dir: 源图片文件夹
        target_root: 目标文件夹根目录
        ratio: 划分比例，2个元素为(train, val)，3个元素为(train, val, test)
        seed: 随机种子，保证结果可复现
    """
    # 检查源文件夹是否存在
    if not os.path.exists(source_dir):
        print(f"源文件夹 {source_dir} 不存在。")
        return

    # 检查比例合法性：必须是2或3个元素，且总和为1
    if len(ratio) not in (2, 3):
        print("ratio参数必须包含2个或3个元素（如(0.8,0.2)或(0.7,0.2,0.1)）")
        return
    if round(sum(ratio), 2) != 1.0:
        print("划分比例总和必须为1（如(0.8,0.2)或(0.7,0.2,0.1)），请重新设置ratio参数。")
        return

    # 确定数据集类型
    dataset_types = ['train', 'val'] if len(ratio) == 2 else ['train', 'val', 'test']
    print(f"将划分数据集为: {dataset_types}，比例: {ratio}")

    # 1. 按类别分组图片（同时提取数字后缀）
    prefix_groups = defaultdict(list)  # key: 类别名, value: [(原文件名, 数字后缀, 文件路径), ...]
    all_classes = set()
    img_suffixes = ('.jpg', '.jpeg', '.png', '.bmp', '.tiff')  # 支持的图片后缀
    # 正则表达式:匹配文件名中的数字后缀
    number_pattern = re.compile(r'(\d+)\.\w+$')

    for filename in os.listdir(source_dir):
        # 过滤非图片文件
        if not filename.lower().endswith(img_suffixes):
            print(f"跳过非图片文件:{filename}")
            continue

        # 提取类别（假设格式为"类别_数字.后缀"）
        class_name = filename.split('_')[0]  # 按'_'分割取前缀作为类别
        all_classes.add(class_name)

        # 提取数字后缀
        match = number_pattern.search(filename)
        if not match:
            print(f"警告:文件名 {filename} 中未找到数字后缀，将跳过该文件")
            continue
        number_suffix = match.group(1)  # 提取数字部分

        file_path = os.path.join(source_dir, filename)
        prefix_groups[class_name].append((filename, number_suffix, file_path))

    if not all_classes:
        print("源文件夹中未找到有效图片文件，无法划分数据集。")
        return
    print(f"共检测到 {len(all_classes)} 个类别，开始划分数据集...")

    # 2. 创建目标文件夹结构
    dataset_roots = {
        dtype: os.path.join(target_root, dtype)
        for dtype in dataset_types
    }

    # 创建根目录
    for root in dataset_roots.values():
        os.makedirs(root, exist_ok=True)

    # 为每个类别创建子文件夹
    for class_name in all_classes:
        for root in dataset_roots.values():
            os.makedirs(os.path.join(root, class_name), exist_ok=True)

    print(f"目标文件夹结构已创建:")
    for dtype, root in dataset_roots.items():
        print(f"- {dtype}: {root}")

    # 3. 保存类别映射文件
    classes = sorted(all_classes)
    classes_to_idx = {cls: i for i, cls in enumerate(classes)}
    idx_to_labels = {i: cls for i, cls in enumerate(classes)}

    mapping_dir = target_root
    with open(os.path.join(mapping_dir, 'classes_to_label.txt'), 'w', encoding='utf-8') as f:
        for cls, idx in classes_to_idx.items():
            f.write(f"{cls},{idx}\n")
    with open(os.path.join(mapping_dir, 'label_to_classes.txt'), 'w', encoding='utf-8') as f:
        for idx, cls in idx_to_labels.items():
            f.write(f"{idx},{cls}\n")
    with open(os.path.join(mapping_dir, 'classes.txt'), 'w', encoding='utf-8') as f:
        for cls in classes:
            f.write(f"{cls}\n")
    print(f"类别映射文件已保存至 {mapping_dir}")

    # 4. 按类别划分并移动图片（核心：严格保持比例）
    random.seed(seed)
    for class_name, files in prefix_groups.items():
        total_imgs = len(files)
        if total_imgs == 0:
            print(f"类别 {class_name} 无有效图片，跳过划分。")
            continue

        # 计算各数据集的数量（核心算法）
        if len(ratio) == 2:
            # 处理 train + val 模式
            train_ratio, val_ratio = ratio

            # 计算基础数量
            base_train = total_imgs * train_ratio
            base_val = total_imgs * val_ratio

            # 初步分配（四舍五入）
            train_num = round(base_train)
            val_num = round(base_val)

            # 调整总和
            if train_num + val_num != total_imgs:
                diff = total_imgs - (train_num + val_num)
                if diff > 0:
                    train_num += diff
                else:
                    train_num += diff

            # 确保val为偶数
            if val_num % 2 != 0:
                val_num += 1 if val_num > 0 else 1
                train_num = total_imgs - val_num

                # 极端情况处理
                if train_num < 0:
                    train_num = 0
                    val_num = total_imgs if total_imgs % 2 == 0 else total_imgs - 1

            dataset_nums = [train_num, val_num]
        else:
            # 处理 train + val + test 模式（核心改进）
            train_ratio, val_ratio, test_ratio = ratio
            total_ratio = sum(ratio)  # 应该为1，但保留作为安全检查

            # 计算基础比例系数
            base_train = total_imgs * train_ratio / total_ratio
            base_val = total_imgs * val_ratio / total_ratio
            base_test = total_imgs * test_ratio / total_ratio

            # 初步分配（四舍五入）
            train_num = round(base_train)
            val_num = round(base_val)
            test_num = round(base_test)

            # 调整总和匹配总数
            current_total = train_num + val_num + test_num
            if current_total != total_imgs:
                diff = total_imgs - current_total
                # 按比例分配差异
                ratio_weights = [train_ratio, val_ratio, test_ratio]
                sorted_indices = sorted(range(3), key=lambda i: ratio_weights[i], reverse=True)

                # 分配差异到权重最大的部分
                for i in sorted_indices:
                    if diff == 0:
                        break
                    if i == 0:  # train
                        train_num += diff
                    elif i == 1:  # val
                        val_num += diff
                    else:  # test
                        test_num += diff
                    diff = total_imgs - (train_num + val_num + test_num)

            # 确保val和test为偶数，同时保持它们的比例关系
            vt_ratio = val_ratio / test_ratio  # 目标比例

            # 计算当前val和test的总和
            total_vt = val_num + test_num

            # 寻找最佳偶数组合
            best_val = val_num
            best_test = test_num
            min_diff = float('inf')

            # 只在小范围内搜索最佳偶数组合（±2范围内）
            for val_adj in [-2, 0, 2]:
                candidate_val = val_num + val_adj
                if candidate_val < 0:
                    continue
                if candidate_val % 2 != 0:
                    continue

                candidate_test = total_vt - candidate_val
                if candidate_test < 0:
                    continue
                if candidate_test % 2 != 0:
                    continue

                # 计算与目标比例的差异
                current_ratio = candidate_val / candidate_test if candidate_test != 0 else float('inf')
                diff = abs(current_ratio - vt_ratio)

                if diff < min_diff:
                    min_diff = diff
                    best_val = candidate_val
                    best_test = candidate_test

            # 更新val和test数量
            val_num, test_num = best_val, best_test
            # 相应调整train数量
            train_num = total_imgs - val_num - test_num

            # 最终安全检查
            if train_num < 0:
                train_num = 0
                total_vt = total_imgs
                # 重新计算val和test
                best_val = round(total_vt * val_ratio / (val_ratio + test_ratio))
                best_val = best_val if best_val % 2 == 0 else best_val + 1
                best_test = total_vt - best_val
                best_test = best_test if best_test % 2 == 0 else best_test + 1
                if best_val + best_test > total_vt:
                    best_val -= 2
                val_num, test_num = best_val, best_test

            dataset_nums = [train_num, val_num, test_num]

        # 随机打乱并划分
        random.shuffle(files)
        dataset_files = []
        start = 0
        for num in dataset_nums:
            end = start + num
            dataset_files.append(files[start:end])
            start = end

        # 定义带重命名的移动函数
        def move_with_original_suffix(file_list, target_root_dir, cls_name):
            """将文件移动到目标文件夹，新文件名 = 原数字后缀 + 原扩展名"""
            target_class_dir = os.path.join(target_root_dir, cls_name)
            for filename, number_suffix, src_path in file_list:
                # 提取原文件扩展名
                _, ext = os.path.splitext(filename)
                # 新文件名:数字后缀 + 扩展名
                new_filename = f"{number_suffix}{ext}"
                target_file_path = os.path.join(target_class_dir, new_filename)

                # 处理可能的文件名重复
                if os.path.exists(target_file_path):
                    print(f"警告:{new_filename} 已存在，将重命名为 {number_suffix}_dup{ext}")
                    new_filename = f"{number_suffix}_dup{ext}"
                    target_file_path = os.path.join(target_class_dir, new_filename)

                try:
                    shutil.copy2(src_path, target_file_path)
                except Exception as e:
                    print(f"处理 {filename} 时出错:{str(e)}")

        # 移动并命名各数据集图片
        for i, dtype in enumerate(dataset_types):
            move_with_original_suffix(dataset_files[i], dataset_roots[dtype], class_name)

        # 打印划分结果
        info_parts = []
        for dtype, files in zip(dataset_types, dataset_files):
            count = len(files)
            even_note = "(偶数)" if (dtype in ['val', 'test'] and count % 2 == 0) else ""
            info_parts.append(f"{dtype}: {count:4} {even_note}")

        print(f"类别 {class_name:10} | 总数: {total_imgs:4} | " + " | ".join(info_parts))

    print(f"\n数据集划分完成！最终结构:")
    for dtype, root in dataset_roots.items():
        print(f"- {dtype}集:{root}/[类别名]/[原数字后缀].jpg（如001.jpg）")
    print(f"- 类别映射:{mapping_dir}/[classes_to_label.txt等]")




if __name__ == '__main__':
# 使用示例
    source_dir = r'D:\Code\2-ZSL\0-data\HOB\morl'  # 替换为实际的源文件夹路径
    target_root = r'D:\Code\2-ZSL\0-data\HOB\morl-data'  # 替换为实际的目标文件夹路径
    split_dataset_flat(source_dir, target_root, (0.7, 0.15,0.15))
